import os
import sys
import requests
import subprocess
import winreg
import shutil
import logging
from pathlib import Path
from PyQt6.QtWidgets import QProgressDialog, QApplication, QMessageBox
from PyQt6.QtCore import Qt

class TesseractInstaller:
    def __init__(self):
        self.base_path = self._get_base_path()
        self.temp_dir = os.path.join(self.base_path, 'temp')
        self.tesseract_version = "5.3.1"
        # 使用可靠的下载源
        self.tesseract_urls = [
            "https://digi.bib.uni-mannheim.de/tesseract/tesseract-ocr-w64-setup-5.3.1.20230401.exe",
            "https://digi.bib.uni-mannheim.de/tesseract/tesseract-ocr-w64-setup-v5.3.1.20230401.exe"
        ]
        
        # 设置日志
        self._setup_logging()
        
    def _setup_logging(self):
        """设置日志记录"""
        log_dir = os.path.join(self.base_path, 'logs')
        os.makedirs(log_dir, exist_ok=True)
        log_file = os.path.join(log_dir, 'installer.log')
        
        logging.basicConfig(
            level=logging.DEBUG,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)

    def _get_base_path(self):
        """获取应用程序基础路径"""
        if getattr(sys, 'frozen', False):
            return os.path.dirname(sys.executable)
        return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))

    def _create_progress_dialog(self, title, label_text):
        """创建进度对话框"""
        dialog = QProgressDialog(label_text, "取消", 0, 100)
        dialog.setWindowTitle(title)
        dialog.setWindowModality(Qt.WindowModality.WindowModal)
        dialog.setMinimumDuration(0)
        return dialog

    def _download_file(self, url, save_path, progress_dialog, max_retries=3):
        """下载文件并显示进度，支持重试"""
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Accept': 'application/octet-stream',
            'Accept-Encoding': 'gzip, deflate, br',
            'Connection': 'keep-alive'
        }
        
        for retry in range(max_retries):
            try:
                self.logger.info(f"开始下载文件: {url} (尝试 {retry + 1}/{max_retries})")
                self.logger.info(f"保存路径: {save_path}")
                
                # 使用session来保持连接
                session = requests.Session()
                response = session.get(url, stream=True, headers=headers, timeout=30)
                response.raise_for_status()
                
                total_size = int(response.headers.get('content-length', 0))
                block_size = 8192  # 增加块大小
                downloaded = 0

                with open(save_path, 'wb') as f:
                    for data in response.iter_content(block_size):
                        if progress_dialog.wasCanceled():
                            self.logger.warning("用户取消了下载")
                            return False
                        downloaded += len(data)
                        f.write(data)
                        progress = int((downloaded / total_size) * 100)
                        progress_dialog.setValue(progress)
                        QApplication.processEvents()

                self.logger.info("文件下载完成")
                return True
                
            except requests.exceptions.RequestException as e:
                self.logger.error(f"下载失败 (尝试 {retry + 1}/{max_retries}): {str(e)}")
                if retry < max_retries - 1:
                    wait_time = (retry + 1) * 5  # 递增等待时间
                    self.logger.info(f"等待{wait_time}秒后重试...")
                    QApplication.processEvents()
                    import time
                    time.sleep(wait_time)
                else:
                    QMessageBox.critical(None, "下载错误", 
                        f"下载文件失败：{str(e)}\n"
                        "请检查网络连接或手动下载安装。\n"
                        "您可以从以下地址手动下载：\n"
                        "https://digi.bib.uni-mannheim.de/tesseract/tesseract-ocr-w64-setup-5.3.1.20230401.exe")
                    return False
            except Exception as e:
                self.logger.error(f"下载过程发生错误: {str(e)}")
                QMessageBox.critical(None, "错误", f"下载过程发生错误：{str(e)}")
                return False

    def _add_to_path(self, path):
        """将路径添加到系统环境变量"""
        try:
            self.logger.info(f"正在添加环境变量: {path}")
            key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, 
                               "Environment", 
                               0, 
                               winreg.KEY_ALL_ACCESS)
            
            # 获取当前Path值
            path_value, _ = winreg.QueryValueEx(key, "Path")
            paths = path_value.split(';')
            
            # 如果路径不在Path中，则添加
            if path not in paths:
                paths.append(path)
                new_path = ';'.join(paths)
                winreg.SetValueEx(key, "Path", 0, winreg.REG_EXPAND_SZ, new_path)
                self.logger.info("环境变量添加成功")
            else:
                self.logger.info("环境变量已存在")
            
            winreg.CloseKey(key)
            return True
        except Exception as e:
            self.logger.error(f"添加环境变量失败: {str(e)}")
            QMessageBox.critical(None, "环境变量错误", f"添加环境变量失败：{str(e)}")
            return False

    def install_tesseract(self):
        """安装Tesseract-OCR"""
        try:
            # 创建临时目录
            os.makedirs(self.temp_dir, exist_ok=True)
            self.logger.info(f"创建临时目录: {self.temp_dir}")
            
            # 创建进度对话框
            progress_dialog = self._create_progress_dialog(
                "安装Tesseract-OCR",
                "正在下载Tesseract-OCR安装程序..."
            )
            
            # 尝试从多个源下载安装程序
            installer_path = os.path.join(self.temp_dir, "tesseract_installer.exe")
            download_success = False
            
            for url in self.tesseract_urls:
                if self._download_file(url, installer_path, progress_dialog):
                    download_success = True
                    break
                else:
                    self.logger.warning(f"从 {url} 下载失败，尝试下一个源")
            
            if not download_success:
                self.logger.error("所有下载源都失败")
                QMessageBox.critical(None, "下载错误", "无法下载Tesseract-OCR安装程序，请检查网络连接或手动下载安装。")
                return False
            
            progress_dialog.setLabelText("正在安装Tesseract-OCR...")
            progress_dialog.setValue(0)
            
            # 运行安装程序
            self.logger.info("开始运行安装程序")
            try:
                # 使用管理员权限运行安装程序，并自动安装英语和日语语言包
                subprocess.run([installer_path, '/S', '/LANG=eng+jpn'], check=True, shell=True)
                self.logger.info("安装程序运行完成")
            except subprocess.CalledProcessError as e:
                self.logger.error(f"安装程序运行失败: {str(e)}")
                QMessageBox.critical(None, "安装错误", f"安装程序运行失败：{str(e)}")
                return False
            
            # 添加环境变量
            tesseract_path = r"C:\Program Files\Tesseract-OCR"
            if not self._add_to_path(tesseract_path):
                return False
            
            return True
        except Exception as e:
            self.logger.error(f"安装过程发生错误: {str(e)}")
            QMessageBox.critical(None, "安装错误", f"安装过程发生错误：{str(e)}")
            return False

    def cleanup(self):
        """清理临时文件"""
        try:
            if os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
                self.logger.info("临时文件清理完成")
        except Exception as e:
            self.logger.error(f"清理临时文件失败: {str(e)}")

    def check_installation(self):
        """检查Tesseract-OCR是否已安装"""
        try:
            # 检查安装路径
            tesseract_path = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
            if not os.path.exists(tesseract_path):
                self.logger.info("Tesseract-OCR未安装")
                return False
            
            self.logger.info("Tesseract-OCR已正确安装")
            return True
        except Exception as e:
            self.logger.error(f"检查安装状态时发生错误: {str(e)}")
            return False

    def _download_language_pack(self, lang_code, save_dir):
        """下载指定语言包到tessdata目录"""
        lang_urls = {
            'eng': 'https://github.com/tesseract-ocr/tessdata/raw/main/eng.traineddata',
            'jpn': 'https://github.com/tesseract-ocr/tessdata/raw/main/jpn.traineddata',
        }
        url = lang_urls.get(lang_code)
        if not url:
            self.logger.error(f"未知语言包: {lang_code}")
            return False
        save_path = os.path.join(save_dir, f'{lang_code}.traineddata')
        if os.path.exists(save_path):
            self.logger.info(f"{lang_code}.traineddata 已存在，无需下载")
            return True
        # 创建进度对话框
        progress_dialog = self._create_progress_dialog(
            f"下载{lang_code}语言包",
            f"正在下载 {lang_code}.traineddata ..."
        )
        result = self._download_file(url, save_path, progress_dialog)
        if result:
            self.logger.info(f"{lang_code}.traineddata 下载完成")
        else:
            self.logger.error(f"{lang_code}.traineddata 下载失败")
        return result

    def _ensure_language_packs(self):
        """确保eng和jpn语言包都存在"""
        tessdata_dir = os.path.join(r"C:\Program Files\Tesseract-OCR", "tessdata")
        os.makedirs(tessdata_dir, exist_ok=True)
        for lang in ["eng", "jpn"]:
            self._download_language_pack(lang, tessdata_dir)

    def install(self):
        """执行完整的安装过程"""
        if self.check_installation():
            self._ensure_language_packs()
            return True
        try:
            self.logger.info("开始安装过程")
            # 安装Tesseract-OCR
            if not self.install_tesseract():
                self.logger.error("Tesseract-OCR安装失败")
                return False
            # 安装语言包
            self._ensure_language_packs()
            # 清理临时文件
            self.cleanup()
            self.logger.info("安装过程完成")
            return True
        except Exception as e:
            self.logger.error(f"安装过程出错: {str(e)}")
            QMessageBox.critical(None, "安装错误", f"安装过程出错：{str(e)}")
            return False 